## Copyright 2002-2004 Andrew Loewenstern, All Rights Reserved
## Copyright 2007 Roee Shlomo, All Rights Reserved
# see LICENSE.txt for license information
import sys
from traceback import print_exc, format_exc
from utility import tobinary16
from defer import Deferred
from BitTornado.bencode import bencode, bdecode
from BitTornado.clock import clock
from const import *
from errors import *

DEBUG  = False
DEBUG2 = False

class NodeBanEntry(object):
    __slots__ = ('count', 'src', 'limit')
    def __init__(self):
        self.count = 0
        self.src = ""
        self.limit = 0

class KRPC:
    """
    KRPC Handler
    """
    def __init__(self, factory):
        self.factory = factory
        self.rawserver = factory.rawserver
        self.tids = {}
        self.mtid = 0
        self.bans = [NodeBanEntry() for i in xrange(20)]
        
    def getTid(self):
        self.mtid += 1
        return tobinary16(self.mtid % 0xFFFF)
        
    def datagramReceived(self, data, addr):
        # Handle node bans - algorithm taken from libtorrent
        match = None
        minimal = self.bans[0]
        now = clock()
        for i in self.bans:
            if i.src == addr[0]:
                match = i
                break
            if i.count < minimal.count:
                minimal = i
        if match is not None:
            match.count += 1
            if match.count >= 20:
                if now < match.limit:
                    #if match.count == 20:
                    #    print "Debug: DHT node banned for 5 minutes %s" % addr[0]
                    match.limit = now + 300
                    return
                match.count = 0 
                match.limit = now + 5
        else:
            minimal.count = 1
            minimal.limit = now + 5
            minimal.src = addr[0]

        # Handle message
        try:
            msg = bdecode(data)
        except:
            if DEBUG:
                print_exc()
        else:
            if DEBUG:
                print 'datagramReceived', msg
            if msg[TYP]  == REQ:
                f = getattr(self.factory ,"krpc_" + msg[REQ], None)
                msg[ARG]['_krpc_sender'] =  addr
                host = addr[0]
                if f and hasattr(f, '__call__'):
                    try:
                        ret = f(**msg[ARG])
                    except KrpcBaseError, e:
                        if DEBUG2:
                            print e, addr, msg.get(VER, "unknown client")
                        self.sendMessage({TID:msg[TID], TYP:ERR, ERR : e.get()}, addr)                            
                    except:
                        print msg.get(VER, "unknown client")
                        print_exc()
                    else:
                        if not ret:
                            ret = {}
                        self.sendMessage({TID : msg[TID], TYP : RSP, RSP : ret}, addr)
                else:
                    self.sendMessage({TID:msg[TID], TYP:ERR, ERR : KrpcMethodUnknownError().get()}, addr)
            elif msg[TYP] == RSP:
                if msg[TID] in self.tids:
                    df = self.tids[msg[TID]]
                    del self.tids[msg[TID]]
                    df.callback({'rsp' : msg[RSP], '_krpc_sender': addr})
                else:
                    pass

            elif msg[TYP] == ERR:
                if DEBUG2:
                    print "Error: DHT - ", msg
                if msg[TID] in self.tids:
                    df = self.tids[msg[TID]]
                    del self.tids[msg[TID]]
                    df.errback(msg[ERR])
                else:
                    pass
                
            else:
                if msg[TID] in self.tids:
                    df = self.tids[msg[TID]]
                    del self.tids[msg[TID]]
                    df.errback(KRPC_ERROR_RECEIVED_UNKNOWN)
                else:
                    pass
                
    def sendRequest(self, method, args, host, port):
        msg = {TID : self.getTid(), TYP : REQ,  REQ : method, ARG : args}
        def timeOut(tids = self.tids, id = msg[TID]):
            if id in tids:
                df = tids[id]
                del tids[id]
                df.errback(KRPC_ERROR_TIMEOUT)
        d = Deferred()
        self.tids[msg[TID]] = d
        self.rawserver.add_task(timeOut, KRPC_TIMEOUT)
        self.sendMessage(msg, (host, port))
        return d

    def sendMessage(self, msg, addr):
        msg[VER] = "LH" + tobinary16(BUILD_NUMBER)
        self.rawserver.write(bencode(msg), addr)
        if DEBUG:
             print 'sendMessage', msg
